package gold.utils;

import java.util.*;

/**
 * Created by fanzhenyu02 on 2020/6/27.
 * Union Find 算法实现
 */
public class UF {
    // 连通分量个数
    private int count;
    // 存储一棵树
    private int[] parent;
    // 记录树的“重量”
    private int[] size;

    public UF(int n) {
        this.count = n;
        parent = new int[n];
        size = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            size[i] = 1;
        }
    }

    /* 将 p 和 q 连接 */
    public void union(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        if (rootP == rootQ)
            return;

        // 小树接到大树下面，较平衡
        if (size[rootP] > size[rootQ]) {
            parent[rootQ] = rootP;
            size[rootP] += size[rootQ];
        } else {
            parent[rootP] = rootQ;
            size[rootQ] += size[rootP];
        }
        count--;
    }

    /* 判断 p 和 q 是否连通 */
    public boolean connected(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        return rootP == rootQ;
    }

    private int find(int x) {
        while (parent[x] != x) {
            // 进行路径压缩
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    }

    /* 返回图中有多少个连通分量 */
    public int count() {
        return count;
    }

    /* 返回图中联通集合 */
    public Map<Integer, List<Integer>> outputAggregateSet(int n) {
        Map<Integer, List<Integer>> aggregateMap = new HashMap<>();

        // 共有n个集合
        for (int i = 0; i < n; i++) {
            int parentId = find(i);
            if (!aggregateMap.containsKey(parentId)) aggregateMap.put(parentId, new ArrayList<>());
            aggregateMap.get(parentId).add(i);
        }

        System.out.println(aggregateMap.toString());
        return aggregateMap;
    }

    public static void main(String[] args) {
        /**
         * Given input of (1,2), (2,3), (4,5), (7,8), (5,6,7)
         * Return (1,2,3), (4,5,6,7,8)
         */

        UF uf = new UF(9);
        uf.union(1, 2);
        uf.union(2, 3);
        uf.union(4, 5);
        uf.union(7, 8);
        uf.union(5, 6);
        uf.union(6, 7);

        final Map<Integer, List<Integer>> integerListMap = uf.outputAggregateSet(9);
        System.out.println(integerListMap.toString());
    }

}
